package com.beautiful.xin.rbtree;

/**
 * @author xuqingxin 2021/7/5 2:58 下午
 */
public class RBTree<K extends Comparable<K>, V> {
    private static final boolean RED = false;
    private static final boolean BLACK = true;

    private RBNode root;

    public RBNode getRoot() {
        return root;
    }

    public void setRoot(RBNode root) {
        this.root = root;
    }

    /**
     * 新增节点
     * 在红黑树节点的新增
     * 1.普通的二叉树插入
     * 先查询插入到的位置
     * 将kv封装为Node对象插入到Tree中
     * 2. 红黑树的平衡(调整  旋转+变色)
     *
     * @param key
     * @param value
     */
    public void put(K key, V value) {
        RBNode t = this.root;
        if (root == null) {
            // 说明红黑树是空的 插入的就是root
            root = new RBNode(key, value == null ? key : value, null);
            return;
        }
        // 查找插入的位置
        int cmp;
        // 记录寻找节点的父节点
        RBNode parent;
        if (key == null) {
            throw new NullPointerException();
        }

        do {
            parent = t;
            cmp = key.compareTo((K) t.key);

            if (cmp < 0) {
                t = t.left;
            } else if (cmp > 0) {
                t = t.right;
            } else {
                t.setValue(value == null ? key : value);
                return;
            }
        } while (t != null);

        // 找到插入位置 parent的子节点
        RBNode e = new RBNode(key, value == null ? key : value, parent);
        if (cmp < 0) {
            parent.left = e;
        } else {
            parent.right = e;
        }

        // 做平衡处理 旋转,变色
        fixAfterPut(e);


    }

    /**
     * 删除操作:
     * 1.删除叶子节点, 直接删除
     * 2.删除的节点有一个子节点, 那么用子节点来替代
     * 3.如果删除的节点有两个子节点,此时需要找到前驱节点或者后继节点来替代 可以转换为1 2的情况
     *
     * <pre>
     *     删除节点有4中情况
     *     A: 删除的是叶子节点且该叶子节点是红色的 ---> 无需修  (删除 P2r 或 P2r')
     *          +------------------------------+
     *          |     P1b       |      P1b     |
     *          |    /   \      |              |
     *          |  P2r   P2r'   |              |
     *          +------------------------------+
     *     B: 删除的是叶子节点且该叶子节点是黑色的 ---> 很明显会破坏特性5，需要修复。
     *
     *     C. 删除的节点（为了便于叙述我们将其称为P）下面有一个子节点 S，对于这种情况我们通过 将P和S的值交换的方式，巧妙的将删除P变为删除S，S是叶子节点，
     *          这样C这种情况就会转 换为A, B这两种情况：
     *          C1： P为黑色，S为红色 ---> 对应 A 这种情况
     *          C2: P为黑色或红色，S为黑色 --- > 对应 B 这种情况
     *
     *     D. 删除的节点有两个子节点，对于这种情况，我们通过将P和它的后继节点N的值交换的方 式，将删除节点P转换为删除后继节点N，而后继节点只可能是以下两种情况：
     *          D1: N是叶子节点 --- > 对应情况 A 或 B
     *          D2: N有一个子节点 ---- > 对应情况 C
     *
     *       所以通过上面的分析我们发现，红黑树节点删除后的修复操作都可以转换为 A 或 B这两种情况，而A不需要修复，所以我们只需要研究B这种情况如何修复就行了。
     *
     *       下面我们讨论如何修复B中情况：
     *  -------------------------------------------------------------------------------------------------------
     *       B:
     *       删除的节点是左节点 P2b 时:
     *       1. 兄弟节点是黑色 , 且有一个左子节点(可以判断为红色)
     *                原始        |  兄弟变红,兄弟儿子变黑   |     根据兄弟右旋          |           删除         |  根据删除节点父节点左旋
     *         +-------------------------------------------------------------------------------------------------------------------+
     *         |       P1x       |         P1x            |         P1x             |          P1x           |         P4x         |
     *         |      /   \      |        /   \           |        /   \            |         /   \          |        /   \        |
     *         |    P2b   P3b    |      P2b   P3r         |      P2b   P4b          |       null   P4b       |      P1b   P3b      |
     *         |         /       |           /            |               \         |                \       |     /               |
     *         |       p4r       |         p4b            |               P3r       |                P3r     |   null              |
     *         +-------------------------------------------------------------------------------------------------------------------+
     *       2. 兄弟节点是黑色 , 且有一个右子节点(可以判断为红色)
     *                原始           |           删除         |  根据删除节点父节点左旋   (父节点颜色给兄弟节点, 父节点和兄弟右子节点变黑)
     *         +-------------------------------------------------------------------------------------------------------------------+
     *         |       P1x          |          P1x           |         P3x                                                         |
     *         |      /   \         |         /   \          |        /   \                                                        |
     *         |    P2b   P3b       |       null   P3b       |      P1b   P4b                                                      |
     *         |             \      |                \       |     /                                                               |
     *         |              p4r   |                P4r     |   null                                                              |
     *         +-------------------------------------------------------------------------------------------------------------------+
     *       3. 兄弟节点是黑色 , 且有左右两个子节点(可以判断为红色)和1 情况相同
     *                原始           |           删除         |  根据删除节点父节点左旋   (父节点颜色给兄弟节点, 父节点和兄弟右子节点变黑)
     *         +-------------------------------------------------------------------------------------------------------------------+
     *         |       P1x          |          P1x           |         P3x                                                         |
     *         |      /   \         |         /   \          |        /   \                                                        |
     *         |    P2b   P3b       |       null   P3b       |      P1b   P4b                                                      |
     *         |          /  \      |             /  \       |     /   \                                                           |
     *         |        P5r   P4r   |           P5r   P4r    |   null   P5r                                                        |
     *         +-------------------------------------------------------------------------------------------------------------------+
     *       4. 兄弟节点是黑色 没有子节点
     *                原始           |           删除    兄弟节点变红,将父节点设置为当前递归节点 直到root或遇到红色节点
     *         +-------------------------------------------------------------------------------------------------------------------+
     *         |       P1x          |          P1b                                                                                 |
     *         |      /   \         |         /   \                                                                                |
     *         |    P2b   P3b       |       null   P3r                                                                             |
     *         |                    |                                                                                              |
     *         |                    |                                                                                              |
     *         +-------------------------------------------------------------------------------------------------------------------+
     *       5. 兄弟节点是红色 , 那么兄弟节点则有两个黑色子节点
     *                原始           | (父节点颜色给兄弟节点, 父节点和兄弟右子节点变黑,左节点变红)   |  根据删除节点父节点左旋
     *         +-------------------------------------------------------------------------------------------------------------------+
     *         |       P1b          |          P1b                                            |         P3b                        |
     *         |      /   \         |         /   \                                           |        /   \                       |
     *         |    P2b   P3r       |       null   P3b                                        |      P1b   P4b                     |
     *         |          /  \      |             /  \                                        |     /   \                          |
     *         |        P5b   P4b   |           P5r   P4b                                     |   null   P5r                       |
     *         +-------------------------------------------------------------------------------------------------------------------+
     *
     *  </pre>
     */

    public V remove(K key) {
        //先根据key 找到对应节点
        RBNode node = getNode(key);
        if (node == null) {
            return null;
        }
        // 把删除节点原来的值 保存下来
        V oldValue = (V) node.value;
        // 删除节点的方法
        deleteNode(node);
        return oldValue;
    }

    /**
     * 删除操作:
     * 1.删除叶子节点, 直接删除
     * 2.删除的节点有一个子节点, 那么用子节点来替代
     * 3.如果删除的节点有两个子节点,此时需要找到前驱节点或者后继节点来替代 可以转换为1 2的情况
     *
     * @param node
     */
    private void deleteNode(RBNode node) {
        // 情况三
        if (leftOf(node) != null && rightOf(node) != null) {
            // 找到要删除节点的后继节点
            RBNode seccessor = seccessor(node);
            // 然后用后继节点的值 覆盖点要删除节点的信息
            node.key = seccessor.key;
            node.value = seccessor.value;
            // 然后我们要删除的节点就变成了后继节点
            node = seccessor;
        }
        // 获取要替换的节点
        RBNode r = node.left != null ? node.left : node.right;
        if (r != null) {

            r.parent = node.parent;

            if (node.parent == null) {
                root = r;
            } else if (node == leftOf(parentOf(node))) {
                parentOf(node).left = r;
            } else {
                parentOf(node).right = r;
            }

            //将node的左右和父都指向null 等待gc
            node.left = node.right = node.parent = null;
            //替换完成后需要调整平衡
            if (node.color == BLACK) {
                fixAfterRemove(node);
            }
        } else if (node.parent == null) {
            root = null;
        } else {
            // 情况1
            if (node.color == BLACK) {
                fixAfterRemove(node);
            }
            // 删除
            if (node.parent != null) {
                if (node == parentOf(node).left) {
                    parentOf(node).left = null;
                } else {
                    parentOf(node).right = null;
                }
                node = null;
            }
        }

    }

    /**
     * 删除后的调整
     *
     * @param x
     */
    private void fixAfterRemove(RBNode x) {

        // 情况2 3
        if (x != root && colorOf(x) == BLACK) {
            if (x == leftOf(parentOf(x))) {
                // 查找兄弟节点
                RBNode r = rightOf(parentOf(x));
                if (colorOf(r) == RED) {
                    setColor(r, BLACK);
                    setColor(parentOf(x), RED);
                    leftRotate(parentOf(x));
                    r = rightOf(parentOf(x));
                }

                if (colorOf(leftOf(r)) == BLACK && colorOf(rightOf(r)) == BLACK) {
                    // 2 兄弟节点没有
                    setColor(r, RED);
                    x = parentOf(x);
                } else {
                    // 1.跟兄弟借 兄弟节点有
                    if (colorOf(rightOf(r)) == BLACK) {
                        setColor(r, RED);
                        setColor(leftOf(r), BLACK);
                        rightRotate(r);
                        r = rightOf(parentOf(x));
                    }
                    setColor(r, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(rightOf(r), BLACK);
                    leftRotate(parentOf(x));
                    x = root;
                }
            } else {
                // 查找兄弟节点
                RBNode r = leftOf(parentOf(x));
                if (colorOf(r) == RED) {
                    setColor(r, BLACK);
                    setColor(parentOf(x), RED);
                    leftRotate(parentOf(x));
                    r = rightOf(parentOf(x));
                }

                if (colorOf(rightOf(r)) == BLACK && colorOf(leftOf(r)) == BLACK) {
                    // 2 兄弟节点没有
                    setColor(r, RED);
                    x = parentOf(x);
                } else {
                    // 1.跟兄弟借 兄弟节点有
                    if (colorOf(leftOf(r)) == BLACK) {
                        setColor(r, RED);
                        setColor(rightOf(r), BLACK);
                        leftRotate(r);
                        r = leftOf(parentOf(x));
                    }
                    setColor(r, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(leftOf(r), BLACK);
                    rightRotate(parentOf(x));
                    x = root;
                }
            }
        }

        //情况1 替代的节点是红色 直接设置为黑色
        setColor(x, BLACK);
    }


    /**
     * 查找节点
     *
     * @param key
     * @return
     */
    private RBNode getNode(K key) {
        RBNode node = this.root;
        while (node != null) {
            int cmp = key.compareTo((K) node.key);
            if (cmp > 0) {
                node = node.right;
            } else if (cmp < 0) {
                node = node.left;
            } else {
                return node;
            }
        }
        return null;
    }


    private boolean colorOf(RBNode node) {
        return node == null ? BLACK : node.color;
    }

    private RBNode parentOf(RBNode node) {
        return node != null ? node.parent : null;
    }

    private RBNode leftOf(RBNode node) {
        return node != null ? node.left : null;
    }

    private RBNode rightOf(RBNode node) {
        return node != null ? node.right : null;
    }

    private void setColor(RBNode node, boolean color) {
        if (node != null) {
            node.setColor(color);
        }
    }

    /**
     * 红黑树的平衡处理
     * <pre>
     * <p>
     *     P--节点
     *     1--节点1
     *     r--Red
     *     b--black
     *
     * 1. 2-3-4树 2节点新增一个元素 直接合并为3节点
     * 红黑树: 新增一个红色节点,这个红色节点会添加在黑色节点下 不需要调整
     *              情况1   |     情况2
     * B  =>        B      |      B
     *                \    |   /
     *                  R  |  R
     *
     *
     * 2. 2-3-4树 3节点新增一个元素 合并为一个4节点
     * 红黑树: 就会有6种情况 两种(左中右)的不需要调整
     * 根左左 根右右 旋转一次
     * 根左右 根右左 旋转两次
     *
     * 红色节点P3r  插入的父亲为爷爷左节点
     * 情况1:
     *         根左    |  P2r左旋 ==>  |   P1b右旋后 ===> |    平衡
     *   +--------------------------------------------------------------+
     *   |    P1b     |      p1b     |        P1b       |       P3b     |
     *   |    /       |      /       |        /         |      /  \     |
     *   |  P2r       |    P2r       |     P3r          |    P2r    P1r |
     *   |            |      \       |    /             |               |
     *   |            |       P3r    |  P2r             |               |
     *   +--------------------------------------------------------------+
     * 情况1 根据父亲左旋后变为 情况2
     * 情况2
     *         根左    |  P1b右旋后 ===>  |    平衡
     *   +-----------------------------------------------+
     *   |    P1b     |        P1b       |       P2b     |
     *   |    /       |        /         |      /  \     |
     *   |  P2r       |     P2r          |    P3r    P1r |
     *   |            |    /             |               |
     *   |            |  P3r             |               |
     *   +-----------------------------------------------+
     *
     * 情况2 根据爷爷右旋后变为 情况3
     * 情况3:插入就平衡
     *         根左    |    平衡
     *   +-----------------------------+
     *   |    P1b     |        P1b     |
     *   |    /       |       /  \     |
     *   |  P2r       |     P2r    P3r |
     *   |            |                |
     *   |            |                |
     *   +-----------------------------+
     *
     *  红色节点P3r  插入的父亲为爷爷右节点
     * 情况1:
     *         根右    |  P2r右旋 ==>  |   P1b左旋后 ===> |    平衡
     *   +---------------------------------------------------------------+
     *   |    P1b     |      p1b     |        P1b        |       P3b     |
     *   |       \    |         \    |           \       |      /  \     |
     *   |       P2r  |         P2r  |           P3r     |    P1r    P2r |
     *   |            |        /     |             \     |               |
     *   |            |      P3r     |             P2r   |               |
     *   +---------------------------------------------------------------+
     * 情况1 根据父亲右旋后变为 情况2
     * 情况2
     *         根右    |  P1b左旋后 ===>  |    平衡
     *   +-----------------------------------------------+
     *   |    P1b     |        P1b       |       P2b     |
     *   |       \    |           \      |      /  \     |
     *   |       P2r  |           P2r    |    P1r    P3r |
     *   |            |             \    |               |
     *   |            |             P3r  |               |
     *   +-----------------------------------------------+
     *
     * 情况2 根据爷爷坐旋后变为 情况3
     * 情况3:插入就平衡
     *         根右    |    平衡
     *   +-----------------------------+
     *   |    P1b     |        P1b     |
     *   |      \     |       /  \     |
     *   |      P2r   |     P3r    P2r |
     *   |            |                |
     *   |            |                |
     *   +-----------------------------+
     *
     * 3. 2-3-4树 4节点 新增一个元素 4节点中间元素升级为父节点 新增元素和剩下的节点合并
     * 红黑树: 新增节点是红色+ 爷爷节点是黑色 父节点和叔叔节点为红色调整为
     * 爷爷节点变为红色, 父亲和叔叔节点变为黑色, 如果爷爷节点为root节点则调整为黑色
     *
     * 红色P4r 插入 爷爷和父亲 叔叔节点变色即可
     *           原始        |     插入后 两红 需变色平衡    |  非Root节点 爷爷变红 后递归  |      Root节点爷爷变黑
     *   +--------------------------------------------------------------------------------------------------------+
     *   |       P1b        |            P1b              |          P1r              |          P1b              |
     *   |      /  \        |         /       \           |       /       \           |       /       \           |
     *   |    P2r    P3r    |      P2r        P3r         |    P2b        P3b         |    P2b        P3b         |
     *   |                  |     /  \       /    \       |   /  \       /    \       |   /  \       /    \       |
     *   |                  |   P4r  P4r'  P4r"   p4r'"   | P4r  P4r'  P4r"   p4r'"   | P4r  P4r'  P4r"   p4r'"   |
     *   +--------------------------------------------------------------------------------------------------------+
     *
     * </pre>
     *
     * @param x
     */
    private void fixAfterPut(RBNode x) {
        // 插入的节点都为红色节点
        setColor(x, RED);
        while (x != null && x != root && parentOf(x).color == RED) {
            // 1. x的父节点是 x爷爷节点的左子节点 4中处理情况
            if (parentOf(x) == parentOf(parentOf(x)).left) {
                // 满足情况的四种情况, 根据是否有叔叔节点 分为两种个情况处理
                // 获取点前节点的叔叔节点
                RBNode y = rightOf(parentOf(parentOf(x)));
                if (colorOf(y) == RED) {
                    // 说明叔叔节点存在, 那就满足情况3
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));

                } else {
                    // 情况2
                    if (x == rightOf(parentOf(x))) {
                        // 插入节点是父节点的右侧节点, 我们需要先根据父节点做一次左旋操作
                        x = parentOf(x);
                        leftRotate(x);
                    }

                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    // 根据爷爷节点做一次右旋操作
                    rightRotate(parentOf(parentOf(x)));
                }
            } else {
                // 与上面情况刚好相反

                // 满足情况的四种情况, 根据是否有叔叔节点 分为两种个情况处理
                // 获取点前节点的叔叔节点
                RBNode y = leftOf(parentOf(parentOf(x)));
                if (colorOf(y) == RED) {
                    // 说明叔叔节点存在, 那就满足情况3
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));

                } else {
                    // 情况2
                    if (x == leftOf(parentOf(x))) {
                        // 插入节点是父节点的右侧节点, 我们需要先根据父节点做一次左旋操作
                        x = parentOf(x);
                        rightRotate(x);
                    }

                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    // 根据爷爷节点做一次右旋操作
                    leftRotate(parentOf(parentOf(x)));
                }

            }


        }

        //根节点都为黑色
        setColor(root, BLACK);
    }


    /**
     * 左旋的实现
     * <pre>
     * +---------------------------------------------+
     * |      p                      pr              |
     * |     /\                     /\               |
     * |    pl pr       ==>        p  pr             |
     * |       /\                 /\                 |
     * |      rl rr             pl  rl               |
     * +---------------------------------------------+
     * </pre>
     * <p>
     * 左旋的时候
     * p-pl 和 pr-rr不变
     * pr-rl 变为 p-rl
     * 判断p节点有没有父节点
     * 如果没有
     * pr为 root节点
     * 如果有
     * r.parent = p.parent
     * 还要设置pr为p.parent的左右子节点
     * <p>
     * 最后设置
     * p-pr 为 pr-p
     *
     * @param p
     */
    private void leftRotate(RBNode p) {

        if (p != null) {
            RBNode r = p.right;

            //1. pr-rl -> p-rl
            p.right = r.left;
            if (r.left != null) {
                r.left.parent = p;
            }

            // 2. 判断 p 是否有父节点
            r.parent = p.parent;
            if (p.parent == null) {
                root = r;
            } else if (p.parent.left == p) {
                p.parent.left = r;
            } else {
                p.parent.right = r;
            }

            // 3. 最后设置 p 为 r 的左节点
            r.left = p;
            p.parent = r;
        }
    }

    private void rightRotate(RBNode p) {

        if (p != null) {
            RBNode r = p.left;

            //1. pr-rl -> p-rl
            p.left = r.right;
            if (r.right != null) {
                r.right.parent = p;
            }

            // 2. 判断 p 是否有父节点
            r.parent = p.parent;
            if (p.parent == null) {
                root = r;
            } else if (p.parent.left == p) {
                p.parent.left = r;
            } else {
                p.parent.right = r;
            }

            // 3. 最后设置 p 为 r 的左节点
            r.left = p;
            p.parent = r;
        }
    }

    /**
     * 前驱节点
     *
     * @param node
     * @return
     */
    private RBNode preducessor(RBNode node) {
        if (node == null) {
            return null;
        } else if (node.left != null) {
            RBNode p = node.left;
            while (p.right != null) {
                p = p.right;
            }
            return p;
        } else {
            //这种情况在删除中是不会存在的 但是就 前驱和后继节点的情况来说还是会存在的
            RBNode p = node.parent;
            RBNode ch = node;
            while (p != null && ch == p.left) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    /**
     * 后继节点
     *
     * @param node
     * @return
     */
    private RBNode seccessor(RBNode node) {
        if (node == null) {
            return null;
        } else if (node.right != null) {
            RBNode p = node.right;
            while (p.left != null) {
                p = p.left;
            }
            return p;
        } else {
            //这种情况在删除中是不会存在的 但是就 前驱和后继节点的情况来说还是会存在的
            RBNode p = node.parent;
            RBNode ch = node;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    // 节点类型
    static class RBNode<K extends Comparable<K>, V> {
        private RBNode parent;
        private RBNode left;
        private RBNode right;

        private boolean color;
        private K key;
        private V value;

        public RBNode(RBNode parent, RBNode left, RBNode right, boolean color, K key, V value) {
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.color = color;
            this.key = key;
            this.value = value;
        }

        public RBNode(K key, V value, RBNode parent) {
            this.parent = parent;
            this.key = key;
            this.value = value;
        }

        public RBNode() {
        }

        public RBNode getParent() {
            return parent;
        }

        public void setParent(RBNode parent) {
            this.parent = parent;
        }

        public RBNode getLeft() {
            return left;
        }

        public void setLeft(RBNode left) {
            this.left = left;
        }

        public RBNode getRight() {
            return right;
        }

        public void setRight(RBNode right) {
            this.right = right;
        }

        public boolean isColor() {
            return color;
        }

        public void setColor(boolean color) {
            this.color = color;
        }

        public K getKey() {
            return key;
        }

        public void setKey(K key) {
            this.key = key;
        }

        public V getValue() {
            return value;
        }

        public void setValue(V value) {
            this.value = value;
        }
    }

}
